# koroFileHeader at Yungoal acer
# Create: 2021-07-30 15:45:06
# LastEdit: 2021-08-02 12:54:23
"""说明暂无"""
__author__ = '749B'


import sys
import textwrap
import unittest

if __name__ == '__main__':
    import os
    sys.path.append(os.getcwd())
# pylint: disable = import-error
from jobs import ip_in_nsg


class IpInNsgStrQueryTestCase(unittest.TestCase):

    def test_str_query_base(self):
        strQuery = ip_in_nsg.str_query_base()
        expect = textwrap.dedent("""\
        Resources
        | where type == "microsoft.network/networksecuritygroups"
        | join kind=leftouter (
            Resourcecontainers
            | where type =~ "microsoft.resources/subscriptions"
            | project SubName=name, subscriptionId
        ) on subscriptionId
        | project nsg_name=name, location, resourceGroup, SubName, properties
        | extend 
            dsr=properties.defaultSecurityRules, 
            sr=properties.securityRules
        | mvexpand asr=array_concat(dsr, sr)
        | project-away dsr, sr, properties
        | extend
            saps=asr.properties.sourceAddressPrefixes,
            sap=asr.properties.sourceAddressPrefix,
            daps=asr.properties.destinationAddressPrefixes,
            dap=asr.properties.destinationAddressPrefix
        | mvexpand ip_addr=array_concat(saps, pack_array(sap), daps, pack_array(dap))
        | project-away saps, sap, daps, dap
        """)
        self.assertEqual(expect, strQuery)

    def test_str_query_base_with_args(self):
        strQuery = ip_in_nsg.str_query_base('3a239097-914b-4a47-9e1a-98c293cd6b7a', 'resource group name')
        expect = textwrap.dedent("""\
        Resources
        | where type == "microsoft.network/networksecuritygroups"
        | where subscriptionId =~ "3a239097-914b-4a47-9e1a-98c293cd6b7a"
        | where resourceGroup =~ "resource group name"
        | join kind=leftouter (
            Resourcecontainers
            | where type =~ "microsoft.resources/subscriptions"
            | project SubName=name, subscriptionId
        ) on subscriptionId
        | project nsg_name=name, location, resourceGroup, SubName, properties
        | extend 
            dsr=properties.defaultSecurityRules, 
            sr=properties.securityRules
        | mvexpand asr=array_concat(dsr, sr)
        | project-away dsr, sr, properties
        | extend
            saps=asr.properties.sourceAddressPrefixes,
            sap=asr.properties.sourceAddressPrefix,
            daps=asr.properties.destinationAddressPrefixes,
            dap=asr.properties.destinationAddressPrefix
        | mvexpand ip_addr=array_concat(saps, pack_array(sap), daps, pack_array(dap))
        | project-away saps, sap, daps, dap
        """)
        self.assertEqual(expect, strQuery)

    def test_str_query_summarize(self):
        base_query = ip_in_nsg.str_query_base()
        strQuery = ip_in_nsg.str_query_summarize(base_query)
        expect = textwrap.dedent("""\
        Resources
        | where type == "microsoft.network/networksecuritygroups"
        | join kind=leftouter (
            Resourcecontainers
            | where type =~ "microsoft.resources/subscriptions"
            | project SubName=name, subscriptionId
        ) on subscriptionId
        | project nsg_name=name, location, resourceGroup, SubName, properties
        | extend 
            dsr=properties.defaultSecurityRules, 
            sr=properties.securityRules
        | mvexpand asr=array_concat(dsr, sr)
        | project-away dsr, sr, properties
        | extend
            saps=asr.properties.sourceAddressPrefixes,
            sap=asr.properties.sourceAddressPrefix,
            daps=asr.properties.destinationAddressPrefixes,
            dap=asr.properties.destinationAddressPrefix
        | mvexpand ip_addr=array_concat(saps, pack_array(sap), daps, pack_array(dap))
        | project-away saps, sap, daps, dap
        | summarize count() by tostring(ip_addr)
        """)
        self.assertEqual(expect, strQuery)

    def test_str_query_one_ip(self):
        base_query = ip_in_nsg.str_query_base()
        ip_addr = '192.168.1.1'
        strQuery = ip_in_nsg.str_query_one_ip(base_query, ip_addr)
        expect = textwrap.dedent("""\
        Resources
        | where type == "microsoft.network/networksecuritygroups"
        | join kind=leftouter (
            Resourcecontainers
            | where type =~ "microsoft.resources/subscriptions"
            | project SubName=name, subscriptionId
        ) on subscriptionId
        | project nsg_name=name, location, resourceGroup, SubName, properties
        | extend 
            dsr=properties.defaultSecurityRules, 
            sr=properties.securityRules
        | mvexpand asr=array_concat(dsr, sr)
        | project-away dsr, sr, properties
        | extend
            saps=asr.properties.sourceAddressPrefixes,
            sap=asr.properties.sourceAddressPrefix,
            daps=asr.properties.destinationAddressPrefixes,
            dap=asr.properties.destinationAddressPrefix
        | mvexpand ip_addr=array_concat(saps, pack_array(sap), daps, pack_array(dap))
        | project-away saps, sap, daps, dap
        | extend
            priority=asr.properties.priority,
            role_name=asr.name,
            access=asr.properties.access,
            dports=asr.properties.destinationPortRanges,
            dport=asr.properties.destinationPortRange,
            sports=asr.properties.sourcePortRanges,
            sport=asr.properties.sourcePortRange
        | where ip_addr !in~ ("", "*")
        | extend ports=set_difference(array_concat(dports, pack_array(dport), sports, pack_array(sport)), pack_array("*"))
        | project-away dports, dport, sports, sport, asr
        | where ip_addr =~ "192.168.1.1"
        """)
        self.assertEqual(expect, strQuery)

    def test_str_query_ips(self):
        base_query = ip_in_nsg.str_query_base()
        ips = ['192.168.100.1', '192.168.100.2', '192.168.100.3']
        strQuery = ip_in_nsg.str_query_ips(base_query, ips)
        expect = textwrap.dedent("""\
        Resources
        | where type == "microsoft.network/networksecuritygroups"
        | join kind=leftouter (
            Resourcecontainers
            | where type =~ "microsoft.resources/subscriptions"
            | project SubName=name, subscriptionId
        ) on subscriptionId
        | project nsg_name=name, location, resourceGroup, SubName, properties
        | extend 
            dsr=properties.defaultSecurityRules, 
            sr=properties.securityRules
        | mvexpand asr=array_concat(dsr, sr)
        | project-away dsr, sr, properties
        | extend
            saps=asr.properties.sourceAddressPrefixes,
            sap=asr.properties.sourceAddressPrefix,
            daps=asr.properties.destinationAddressPrefixes,
            dap=asr.properties.destinationAddressPrefix
        | mvexpand ip_addr=array_concat(saps, pack_array(sap), daps, pack_array(dap))
        | project-away saps, sap, daps, dap
        | extend
            priority=asr.properties.priority,
            role_name=asr.name,
            access=asr.properties.access,
            dports=asr.properties.destinationPortRanges,
            dport=asr.properties.destinationPortRange,
            sports=asr.properties.sourcePortRanges,
            sport=asr.properties.sourcePortRange
        | where ip_addr in~ ("192.168.100.1", "192.168.100.2", "192.168.100.3")
        | extend ports=set_difference(array_concat(dports, pack_array(dport), sports, pack_array(sport)), pack_array("*"))
        | project-away dports, dport, sports, sport, asr
        """)
        self.assertEqual(expect, strQuery)


class IpInNsgTransformIpTestCase(unittest.TestCase):

    def test_ip_addr_only(self):
        ip_list = [
            ('1.1.1.1', ['1.1.1.1', '1.1.1.1/32']),
            ('192.168.1.1', ['192.168.1.1', '192.168.1.1/32']),
            ('172.16.0.1', ['172.16.0.1', '172.16.0.1/32']),
            ('10.20.30.1', ['10.20.30.1', '10.20.30.1/32']),
        ]
        for i in range(4):
            with self.subTest(i=i):
                ip = ip_list[i][0]
                res = ip_list[i][1]
                self.assertListEqual(ip_in_nsg.transform_ip(ip), res)

    def test_ip_addr_with_32(self):
        ip_list = [
            ('1.1.1.2/32', ['1.1.1.2', '1.1.1.2/32']),
            ('192.168.1.2/32', ['192.168.1.2', '192.168.1.2/32']),
            ('172.16.0.2/32', ['172.16.0.2', '172.16.0.2/32']),
            ('10.20.30.2/32', ['10.20.30.2', '10.20.30.2/32']),
        ]
        for i in range(4):
            with self.subTest(i=i):
                ip = ip_list[i][0]
                res = ip_list[i][1]
                self.assertListEqual(ip_in_nsg.transform_ip(ip), res)

    def test_ip_addr_with_net(self):
        ip_list = [
            ('1.1.1.3/8', ['1.1.1.3/8']),
            ('192.168.1.3/30', ['192.168.1.3/30']),
            ('172.16.0.3/24', ['172.16.0.3/24']),
            ('10.20.30.3/16', ['10.20.30.3/16']),
        ]
        for i in range(4):
            with self.subTest(i=i):
                ip = ip_list[i][0]
                res = ip_list[i][1]
                self.assertListEqual(ip_in_nsg.transform_ip(ip), res)
                
    def test_ip_addr_more_case(self):
        ip_list = [
            ('211.144.209.58', ['211.144.209.58', '211.144.209.58/32']),
            ('114.114.114.114/32', ['114.114.114.114', '114.114.114.114/32']),
            # 基本只判断长度，不考虑是否是正确的IP地址
            ('not an ip', ['not an ip', 'not an ip/32']),  # 长度够，就加上网络地址
            ('1/32', ['1/32']),   # 长度不够，就原样返回
            ('2/24', ['2/24']),
            ('not an ip/24', ['not an ip/24']),
            ('not an ip/32', ['not an ip', 'not an ip/32']),
        ]
        for i in range(4):
            with self.subTest(i=i):
                ip = ip_list[i][0]
                res = ip_list[i][1]
                self.assertListEqual(ip_in_nsg.transform_ip(ip), res)


if __name__ == '__main__':
    unittest.main()

